//! Types for describing type layouts.

use std::collections::*;

use super::Error;
use crate::frontend::{Entry, FieldType, ModuleDef, StructDef, StructEntry};
use crate::tysys::UserTypeId;

#[derive(Clone, Debug)]
pub struct ModuleLayout {
    tys: Vec<(String, UserTypeId, StencilLayout)>,
}

impl ModuleLayout {
    pub fn new(tys: Vec<(String, UserTypeId, StencilLayout)>) -> Self {
        Self { tys }
    }

    pub fn iter_tys(&self) -> impl Iterator<Item = &(String, UserTypeId, StencilLayout)> {
        self.tys.iter()
    }

    pub fn find_layout_by_name(&self, s: &str) -> Option<&StencilLayout> {
        self.tys
            .iter()
            .find(|(name, _, _)| name == s)
            .map(|(_, _, l)| l)
    }

    pub fn find_layout_by_id(&self, id: UserTypeId) -> Option<&StencilLayout> {
        self.tys
            .iter()
            .find(|(_, tid, _)| *tid == id)
            .map(|(_, _, l)| l)
    }

    pub fn find_id_by_name(&self, s: &str) -> Option<&UserTypeId> {
        self.tys
            .iter()
            .find(|(name, _, _)| name == s)
            .map(|(_, id, _)| id)
    }

    /// If the type ID exists and corresponds to a struct, returns its size.
    fn get_struct_type_size(&self, id: UserTypeId) -> Option<usize> {
        self.find_layout_by_id(id).and_then(|sl| match sl {
            StencilLayout::Struct(sl) => Some(sl.size),
        })
    }

    /// Computes the size of a field type, recursively.  Returns `None` is
    /// there's a type referenced that doesn't exist.
    pub fn compute_field_type_size(&self, fty: &FieldType) -> Option<usize> {
        let size = match fty {
            FieldType::Void => 0,
            FieldType::Bool => 1,
            FieldType::Int(ity) => ity.byte_len(),
            FieldType::Buf(size) => *size,
            FieldType::Struct(id) => self.get_struct_type_size(*id)?,
            FieldType::Array(size, aty) => size * self.compute_field_type_size(&aty)?,
        };

        Some(size)
    }
}

/// Any of the many kinds of layouts.
#[derive(Clone, Debug)]
pub enum StencilLayout {
    Struct(StructLayout),
}

/// Structs are flat and compact and always have a fixed size.  Fields are
/// usually at fixed offsets within the struct and they can only contain other
/// structs.
#[derive(Clone, Debug)]
pub struct StructLayout {
    fields: Vec<(String, Field)>,
    size: usize,
}

impl StructLayout {
    pub fn new(fields: Vec<(String, Field)>, size: usize) -> Self {
        Self { fields, size }
    }

    pub fn fields(&self) -> &[(String, Field)] {
        &self.fields
    }

    pub fn size(&self) -> usize {
        self.size
    }

    pub fn get_field(&self, name: &str) -> Option<&Field> {
        self.fields.iter().find(|(n, _)| n == name).map(|(_, f)| f)
    }
}

#[derive(Clone, Debug)]
pub struct Field {
    offset: FieldOffset,
    ty: FieldType,
}

impl Field {
    pub fn new(offset: FieldOffset, ty: FieldType) -> Self {
        Self { offset, ty }
    }

    pub fn offset(&self) -> &FieldOffset {
        &self.offset
    }

    pub fn ty(&self) -> &FieldType {
        &self.ty
    }
}

#[derive(Clone, Debug)]
pub enum FieldOffset {
    /// Field at some offset.
    Fixed(u64),
}

/// Records are slightly more complicated than structs, but can have fields
/// that are variable sizes, making their overall size variable.
#[derive(Clone, Debug)]
pub struct RecordLayout {
    fields: Vec<(String, Field)>,
}

/// Context we keep around while generating a module layout.
pub struct ModLayoutCtx {
    struct_sizes: HashMap<UserTypeId, usize>,
}

impl ModLayoutCtx {
    pub fn new() -> Self {
        Self {
            struct_sizes: HashMap::new(),
        }
    }

    pub fn set_size(&mut self, id: UserTypeId, size: usize) {
        self.struct_sizes.insert(id, size);
    }

    pub fn get_size(&self, id: UserTypeId) -> Option<usize> {
        self.struct_sizes.get(&id).copied()
    }
}

/// Generates a module layout from a module def.
pub fn gen_mod_layout_from_mod_defn(mdef: &ModuleDef) -> Result<ModuleLayout, Error> {
    let mut tys = Vec::new();
    let mut mlc = ModLayoutCtx::new();

    for (_, (ename, ent)) in mdef.entries().iter().enumerate() {
        match ent {
            Entry::Struct(id, sdef) => {
                let sl = gen_struct_layout_from_struct_def(sdef, mdef, &mlc)?;
                mlc.set_size(*id, sl.size);
                tys.push((ename.clone(), *id, StencilLayout::Struct(sl)));
            }
        }
    }

    Ok(ModuleLayout::new(tys))
}

fn gen_struct_layout_from_struct_def(
    sdef: &StructDef,
    mdef: &ModuleDef,
    mlc: &ModLayoutCtx,
) -> Result<StructLayout, Error> {
    let mut fields = Vec::new();
    let mut cur_off = 0usize;

    for (_i, e) in sdef.entries().iter().enumerate() {
        match e {
            StructEntry::Field(fd) => {
                let size = get_field_ty_size(fd.ty(), mlc)?;
                let f = Field {
                    offset: FieldOffset::Fixed(cur_off as u64),
                    ty: fd.ty().clone(),
                };
                fields.push((fd.name().to_owned(), f));
                cur_off += size;
            }

            StructEntry::Pad(v) => {
                cur_off += *v as usize;
            }

            StructEntry::PadTo(off) => {
                if *off < cur_off {
                    return Err(Error::PadBeforePos(cur_off, *off));
                }
                cur_off = *off;
            }
        }
    }

    Ok(StructLayout::new(fields, cur_off))
}

fn get_field_ty_size(ty: &FieldType, mlc: &ModLayoutCtx) -> Result<usize, Error> {
    let size = match ty {
        FieldType::Void => 0,
        FieldType::Bool => 1,
        FieldType::Int(ity) => ity.byte_len(),
        FieldType::Buf(size) => *size,
        FieldType::Struct(id) => mlc.get_size(*id).ok_or(Error::StructIdRefBeforeDef(*id))?,
        FieldType::Array(size, aty) => size * get_field_ty_size(aty, mlc)?,
    };

    Ok(size)
}
